from skorch.core import *
from skorch.vision import *
from skorch.vision.io.dcm import read_dcm_info, read_dcm, convert_dcm_grayscale8, BODY_PART_NAME, CHEST_TAG_NAME
import logging as logger
import cv2
from skimage import measure
from skimage import filters
import csv



def scanAllImages(folderPath):
    extensions = ['.dcm']

    images = []

    for root, dirs, files in os.walk(folderPath):
        for file in files:
            if file.lower().endswith(tuple(extensions)):
                relativePath = os.path.join(root, file)
                images.append(relativePath)
    images.sort(key=lambda x: x.lower())
    return images

def brightness(image:NPImage, refer:int=128):
    mean = cv2.mean(image)[0]
    mean = mean-refer
    md = np.abs(image-refer-mean)
    mean_dev = np.mean(md)
    return mean_dev, mean


def brightness_status(mean_dev, mean):
    if mean_dev < abs(mean):
        if mean > 0:
            return 1
        elif mean < 0:
            return 2
        else:
            return 0
    else:
        return 0

def brightness_detect(image_path:Path, refer:int=128, out:Path=''):
    image_info  = read_CR_dcm(str(image_path))
    default_wc = image_info['wc']
    image = convert_dcm_grayscale8(image_info)
    if image.max() == image.min():
        return
    show_image = np.copy(image)
    mean_dev, mean = brightness(image, refer)
    status = brightness_status(mean_dev, mean)

    if status == 1:
        #print(image_path.name, ' 过亮 ', mean, mean_dev)
        image_info['wc'] = default_wc*0.8
        image = convert_dcm_grayscale8(image_info)
        mean_dev, mean = brightness(image, refer)
        status = brightness_status(mean_dev, mean)
        if status == 1:
            print(image_path.name, ' 过亮2 ', mean, mean_dev)
            image_info['wc'] = default_wc * 1.1
            image = convert_dcm_grayscale8(image_info)
            mean_dev, mean = brightness(image, refer)
            status = brightness_status(mean_dev, mean)
        out_name = out.joinpath(image_path.stem + '.png')
        cv2.imwrite(str(out_name), np.hstack((show_image, image)))

    elif status == 2:
        print(image_path.name, ' 过暗 ', mean, mean_dev)
    else:
        print(image_path.name, ' 正常 ')



def clavicle_angle(mask_name):
    angles = {}
    angles['left angle'] = -1
    angles['right angle'] = -1
    image = cv2.imread(mask_name, flags=cv2.IMREAD_GRAYSCALE)
    if image.min() == image.max():
        print(mask_name)
        return angles

    #image = cv2.resize(image, (512,512), interpolation=cv2.INTER_CUBIC)
    thresh = filters.threshold_otsu(image)
    mask = image > thresh

    mask_label = measure.label(mask)
    regionprops = measure.regionprops(mask_label)

    if len(regionprops) > 2:
        print('clavicle_angle regionprops error', mask_name)

    for region in regionprops:
        if region.area > 2000:
            miny, minx, maxy, maxx = region.bbox
            w = maxx - minx
            h = maxy - miny

            angle = np.arctan(float(h) / float(w))

            if minx + w / 2 < image.shape[0] / 2:
                logger.info("left angle: %f" % (angle * 360 / 2 / np.pi))
                an = angle * 360 / 2 / np.pi
                angles['left angle']=round(an, 4)

            if minx + w / 2 > image.shape[0] / 2:
                logger.info("right angle: %f" % (angle * 360 / 2 / np.pi))
                an = angle * 360 / 2 / np.pi
                angles['right angle'] = round(an, 4)
    return angles


def coincidence_rate(lung_name, shoulderblade_image, debug=False):
    """
    计算肺野和肩胛骨的重叠面积
    :param mask1: 肺野二值mask图像
    :param mask2: 肩胛骨二值mask图像
    :param debug:
    :return:
    """
    areas = {}

    lung_image = cv2.imread(lung_name, flags=0)
    if lung_image.min() == lung_image.max():
        areas['left overlap area'] = format(0.0, '0.4f')
        areas['right overlap area'] = format(0.0, '0.4f')
        return areas

    shoulderblade_image = cv2.imread(shoulderblade_image, flags=0)

    if shoulderblade_image.min() == shoulderblade_image.max():
        areas['left overlap area'] = format(0.0, '0.4f')
        areas['right overlap area'] = format(0.0, '0.4f')
        return areas

    thresh = filters.threshold_otsu(lung_image)
    lung_mask = lung_image > thresh

    thresh = filters.threshold_otsu(shoulderblade_image)
    shoulderblade_mask = shoulderblade_image > thresh

    width, height = lung_image.shape

    coin_mask = np.bitwise_and(lung_mask, shoulderblade_mask)

    if debug:
        plt.title('coincidence_image')
        plt.imshow(coin_mask.astype(np.uint8), cmap=plt.cm.gray)
        plt.show()

    coin_mask_label = measure.label(coin_mask)
    coin_regionprops = measure.regionprops(coin_mask_label)

    # if len(coin_regionprops) > 2:
    #     print(lung_name)
    #     plt.title('coincidence_image')
    #     plt.imshow(coin_mask.astype(np.uint8), cmap=plt.cm.gray)
    #     plt.show()
    #     raise('coin_mask regionprops error')

    coin_area = [None, None]
    for region in coin_regionprops:
        if region.area > 2000:
            miny, minx, maxy, maxx  = region.bbox
            if maxx + minx / 2 < width / 2:
                if coin_area[0]:
                    coin_area[0] += region.area #交集图像断裂
                    # print(lung_name)
                    # plt.title('coincidence_image')
                    # plt.imshow(coin_mask.astype(np.uint8), cmap=plt.cm.gray)
                    # plt.show()
                    # raise ('coin_mask left regionprops error')
                else:
                    coin_area[0] = region.area#left
            else:
                if coin_area[1]:
                    coin_area[1] += region.area  # right
                else:
                    coin_area[1] = region.area #right

    shoulderblade_label = measure.label(shoulderblade_mask)
    shoulderblade_regionprops = measure.regionprops(shoulderblade_label)
    shoulderblade_area=[None, None]
    for region in shoulderblade_regionprops:
        miny, minx, maxy, maxx = region.bbox
        if maxx + minx / 2 < width / 2:
            shoulderblade_area[0] = region.area
        else:
            shoulderblade_area[1] = region.area


    for i, (a, b) in enumerate(zip(coin_area, shoulderblade_area)):
        if a and b:
            if i == 0:
                areas['left overlap area'] = round(a/b, 4)
            else:
                areas['right overlap area'] = round(a/b, 4)
        else:
            if i == 0:
                areas['left overlap area'] = format(0.0, '0.4f')
            else:
                areas['right overlap area'] = format(0.0, '0.4f')
    return areas


def process_dr_diocm_josn():
    image_dir = Path('/home/blake/data/dataset/datasets/DR/dicom/images')
    mask_dir = Path('/home/blake/data/dataset/datasets/DR/dicom/masks')


    class_names = ('__ignore__', 'lung', 'clavicle', 'shoulderblade')
    class_name_to_id = {'__ignore__': 0,
                        'lung': 1,
                        'clavicle': 2,
                        'shoulderblade': 3}

    in_file_names = sorted(image_dir.glob('*.json'))

    headers = ['name', 'left angle', 'right angle', 'left overlap area', 'right overlap area']

    with open('stats.csv', 'w', newline='') as f:
        writer = csv.DictWriter(f, headers)
        writer.writeheader()

        for json_file in tqdm(in_file_names):
            out_png_file = mask_dir.joinpath(json_file.with_suffix('.png').name)
            lung_mask_name = out_png_file.parents[0].joinpath('lung', out_png_file.name)
            clavicle_mask_name = out_png_file.parents[0].joinpath('clavicle', out_png_file.name)
            shoulderblade_mask_name = out_png_file.parents[0].joinpath('shoulderblade', out_png_file.name)

            angles_dict = clavicle_angle(str(clavicle_mask_name))
            areas_ratio_dict = coincidence_rate(str(lung_mask_name), str(shoulderblade_mask_name))

            # print(json_file)
            # print(angles_dict, areas_ratio_dict, '\n\n')

            write_dict = {'name':json_file.stem}
            write_dict.update(angles_dict)
            write_dict.update(areas_ratio_dict)
            writer.writerow(write_dict)


def read_dcm_wc(dcm_file):
    dcm_info = read_dcm_info(str(dcm_file))
    return dcm_info['ww'], dcm_info['wc']


def stats_dr_dicom():
    image_dir = [#Path('/home/blake/data/dataset/datasets/DR/dicom/DR/images'),
                 Path('/home/blake/data/dataset/datasets/DR/dicom/DR2/images')]
    in_file_names = []

    for cur_dir in image_dir:
        in_file_names.extend(sorted(cur_dir.glob('*.dcm')))

    ww = []
    wc = []
    # for dcm_file in tqdm(in_file_names):
    #     dcm_info = read_CR_dcm(str(dcm_file))
    #     ww.append(dcm_info['ww'])
    #     wc.append(dcm_info['wc'])

    pool = Pool(16)
    wcs = pool.map(partial(read_dcm_wc), in_file_names)
    pool.close()
    pool.join()

    wcs = np.array(wcs)
    ww = wcs[:,0]
    wc = wcs[:,1]
    np.save('ww', ww)
    np.save('wc', wc)

    print('ww mean: %f, min:%f, max:%f', ww.mean(), ww.max(), ww.min())
    print('wc mean: %f, min:%f, max:%f', wc.mean(), wc.max(), wc.min())



def preprocess_dr_brightness():
    image_dir = Path('/home/blake/data/dataset/datasets/DR/dicom/cache/data/zigong')
    out_dir =  Path('/home/blake/data/dataset/datasets/DR/dicom/show2')

    in_file_names = scanAllImages(str(image_dir))#sorted(image_dir.glob('*.dcm'))

    for dcm_file in (in_file_names):
        #dcm_file = image_dir.joinpath('0','2646-0000004400-3867-10001-1.2.840.887072.1.9.1.1.20180201091510.33.192168207100.dcm')
        brightness_detect(Path(dcm_file), refer=128, out=out_dir)


def move_file(dcm_file, out_dir=''):
    dcm_file = Path(dcm_file)
    image_info = read_CR_dcm(str(dcm_file))
    body = image_info['body_part']
    flag = False

    for key in CHEST_TAG_NAME:
        if key in body:
            flag = True
            break
    #if body == 'CHEST':
    if flag:
        shutil.move(str(dcm_file), str(out_dir))
        if dcm_file.with_suffix('.json').exists():
            shutil.move(str(dcm_file.with_suffix('.json')), str(out_dir))

def preprocess_dr_body_tag():
    image_dir = Path('/home/blake/data/dataset/datasets/DR/data_pool/background')
    out_dir = Path('/home/blake/data/dataset/datasets/DR/data_pool/background_chest')
    # dcm_file = '/home/blake/data/dataset/datasets/DR/data_pool/background/00059191-000000NULL-000NULL-10001-1.2.2018050812215700059191.dcm'
    # image_info = read_CR_dcm(str(dcm_file))

    in_file_names = scanAllImages(str(image_dir))#sorted(image_dir.glob('*.dcm'))

    #for dcm_file in tqdm(in_file_names):
    pool = Pool(16)
    pool.map(partial(move_file, out_dir=out_dir), in_file_names)
    pool.close()
    pool.join()


def read_image_wh(file):
    image = load_image_obj(str(file), loader='opencv')
    return image.shape[0], image.shape[1]

def stats_dr_dicom_wh():
    image_dir = [#Path('/home/blake/data/dataset/datasets/DR/dicom/DR/images'),
                 Path('/home/blake/data/dataset/datasets/DR/dicom/DR2/masks/lung')]
    in_file_names = []

    for cur_dir in image_dir:
        in_file_names.extend(sorted(cur_dir.glob('*.png'))[0:10])

    pool = Pool(8)
    hws = pool.map(partial(read_image_wh), in_file_names)
    pool.close()
    pool.join()

    hws = np.array(hws)
    h = hws[:,0]
    w = hws[:,1]
    print(h.mean(), w.mean())

    fig, axes = plt.subplots(1, 2, figsize=(50, 50))
    axes[0].plot(h)
    axes[0].set_title('DR heights')
    axes[1].plot(h)
    axes[1].set_title('DR widths')
    plt.show()


def mutil_image_mean(path):
    image = load_image_obj(str(path), 'opencv')
    if len(image.shape) < 3:
        image = grey2rgb(image)
    else:
        image = bgr2rgb(image)
    return image_mean(image)


def image_datasets_mean(data_ids, njob=8):
    pool = Pool(njob)
    mean_lists = pool.map(mutil_image_mean, data_ids)
    pool.close()
    pool.join()

    rgbMeans = np.zeros(3)
    for i in range(len(mean_lists)):
        rgbMeans += mean_lists[i]
    return rgbMeans/len(mean_lists)


def mutil_image_std(path):
    image = load_image_obj(str(path), 'opencv')
    if len(image.shape) < 3:
        image = grey2rgb(image)
    else:
        image = bgr2rgb(image)
    return image_std(image)

def image_datasets_std(data_ids, njob=8):
    pool = Pool(njob)
    std_lists = pool.map(mutil_image_std, data_ids)
    pool.close()
    pool.join()

    rgbStds = np.zeros(3)
    for i in range(len(std_lists)):
        rgbStds += std_lists[i]
    return rgbStds/len(std_lists)


def stats_mean_std():
    root = Path("/home/blake/data/dataset/datasets/DR/dicom/DR2")
    images_path = root.joinpath('images')
    masks_path = root.joinpath('masks')

    in_file_names = sorted(images_path.glob('*.png'))

    mean = image_datasets_mean(in_file_names)
    std = image_datasets_std(in_file_names)
    print(mean, std)





if __name__ == '__main__':
    print(" ")
    #process_dr_diocm_josn()
    #stats_dr_dicom()
    #preprocess_dr_brightness()
    #preprocess_dr_body_tag()
    #stats_dr_dicom_wh()
    stats_mean_std()
